import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

"""Non-blocking point-to-point communication."""
#异步通信的核心特征 ： 异步通信的关键在于操作不会立即阻塞当前进程。
# 调用 dist.isend 或 dist.irecv 后，函数会 立即返回一个请求对象（ req ） ，而不是等待通信完成。
# 进程可以在通信过程中继续执行其他任务（例如计算、预处理等），直到显式调用 req.wait() 才会阻塞并等待通信完成。

# 异步通信（如 dist.isend / dist.irecv ）的核心是 非阻塞发起通信 ，但通信的完成需要显式等待（ req.wait() ）。在用户提供的 async_p2p.py 中：

# - dist.isend （异步发送）会立即返回一个请求对象 req ，但数据并未立即发送完成；
# - dist.irecv （异步接收）同理，不会立即阻塞，但数据也未立即接收完成；
# - 必须调用 req.wait() 显式等待通信完成，此时数据才会被完全传输。

# 在 async_p2p.py 中，虽然通信是异步发起的，但最终通过 req.wait() 强制等待通信完成，因此：

# - 进程 0 的 tensor （值为 1 ）会被完整发送到进程 1；
# - 进程 1 接收的 tensor 一定会是进程 0 发送的 1 ，而非初始的 0 ；
# - 最终打印的 tensor[0] 是通信完成后的值，不会出现“未完成传输导致的数据不准确”。
def run(rank, size):
    tensor = torch.zeros(1)
    req = None
    if rank == 0:
        tensor += 1
        # Send the tensor to process 1
        req = dist.isend(tensor=tensor, dst=1)
        print('Rank 0 started sending')
    else:
        # Receive tensor from process 0
        req = dist.irecv(tensor=tensor, src=0)
        print('Rank 1 started receiving')
    req.wait()
    print('Rank ', rank, ' has data ', tensor[0])

def init_process(rank, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)


if __name__ == "__main__":
    size = 2
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
